import sqlite3
import os
import hashlib
import time
from datetime import datetime

# 数据库文件路径
DB_PATH = 'data/pedestrian_tracking.db'

def get_db_connection():
    """创建数据库连接"""
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    """初始化数据库"""
    # 确保data目录存在
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    
    conn = get_db_connection()
    cur = conn.cursor()
    
    # 创建用户表
    cur.execute('''
    CREATE TABLE IF NOT EXISTS users (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        username TEXT UNIQUE NOT NULL,
        password_hash TEXT NOT NULL,
        is_admin BOOLEAN NOT NULL DEFAULT 0,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    ''')
    
    # 创建视频表
    cur.execute('''
    CREATE TABLE IF NOT EXISTS videos (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        filename TEXT NOT NULL,
        original_path TEXT NOT NULL,
        processed_path TEXT,
        uploaded_by INTEGER,
        upload_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        processed BOOLEAN DEFAULT 0,
        duration REAL,
        FOREIGN KEY (uploaded_by) REFERENCES users (id)
    )
    ''')
    
    # 创建轨迹数据表
    cur.execute('''
    CREATE TABLE IF NOT EXISTS trajectories (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        video_id INTEGER NOT NULL,
        track_id INTEGER NOT NULL,
        frame INTEGER NOT NULL,
        x REAL NOT NULL,
        y REAL NOT NULL,
        width REAL NOT NULL,
        height REAL NOT NULL,
        confidence REAL NOT NULL,
        timestamp REAL,
        FOREIGN KEY (video_id) REFERENCES videos (id)
    )
    ''')
    
    # 检查是否已有管理员用户，如果没有则创建默认管理员
    cur.execute("SELECT COUNT(*) FROM users WHERE is_admin = 1")
    if cur.fetchone()[0] == 0:
        # 创建默认管理员账户 (admin/admin)
        password_hash = hashlib.sha256("admin".encode()).hexdigest()
        cur.execute(
            "INSERT INTO users (username, password_hash, is_admin) VALUES (?, ?, ?)",
            ("admin", password_hash, True)
        )
        
        # 创建默认普通用户账户 (user/user)
        password_hash = hashlib.sha256("user".encode()).hexdigest()
        cur.execute(
            "INSERT INTO users (username, password_hash, is_admin) VALUES (?, ?, ?)",
            ("user", password_hash, False)
        )
    
    conn.commit()
    conn.close()

def add_user(username, password, is_admin=False):
    """添加新用户"""
    conn = get_db_connection()
    cur = conn.cursor()
    password_hash = hashlib.sha256(password.encode()).hexdigest()
    
    try:
        cur.execute(
            "INSERT INTO users (username, password_hash, is_admin) VALUES (?, ?, ?)",
            (username, password_hash, is_admin)
        )
        conn.commit()
        success = True
    except sqlite3.IntegrityError:
        success = False
    finally:
        conn.close()
    
    return success

def update_user(user_id, new_username=None, new_password=None, is_admin=None):
    """更新用户信息"""
    conn = get_db_connection()
    cur = conn.cursor()
    
    updates = []
    params = []
    
    if new_username:
        updates.append("username = ?")
        params.append(new_username)
    
    if new_password:
        password_hash = hashlib.sha256(new_password.encode()).hexdigest()
        updates.append("password_hash = ?")
        params.append(password_hash)
    
    if is_admin is not None:
        updates.append("is_admin = ?")
        params.append(is_admin)
    
    if not updates:
        conn.close()
        return False
    
    params.append(user_id)
    query = f"UPDATE users SET {', '.join(updates)} WHERE id = ?"
    
    try:
        cur.execute(query, params)
        conn.commit()
        success = True
    except:
        success = False
    finally:
        conn.close()
    
    return success

def get_all_users():
    """获取所有用户信息"""
    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute("SELECT id, username, is_admin, created_at FROM users")
    users = [dict(row) for row in cur.fetchall()]
    conn.close()
    return users

def delete_user(user_id):
    """删除用户"""
    conn = get_db_connection()
    cur = conn.cursor()
    
    try:
        cur.execute("DELETE FROM users WHERE id = ?", (user_id,))
        conn.commit()
        success = True
    except:
        success = False
    finally:
        conn.close()
    
    return success

def add_video(filename, original_path, user_id):
    """添加视频记录"""
    conn = get_db_connection()
    cur = conn.cursor()
    
    try:
        cur.execute(
            "INSERT INTO videos (filename, original_path, uploaded_by) VALUES (?, ?, ?)",
            (filename, original_path, user_id)
        )
        conn.commit()
        # 获取插入的ID
        video_id = cur.lastrowid
    except:
        video_id = None
    finally:
        conn.close()
    
    return video_id

def update_video_processed(video_id, processed_path, duration):
    """更新视频处理状态"""
    conn = get_db_connection()
    cur = conn.cursor()
    
    try:
        cur.execute(
            "UPDATE videos SET processed = 1, processed_path = ?, duration = ? WHERE id = ?",
            (processed_path, duration, video_id)
        )
        conn.commit()
        success = True
    except:
        success = False
    finally:
        conn.close()
    
    return success

def get_videos(user_id=None):
    """获取视频列表"""
    conn = get_db_connection()
    cur = conn.cursor()
    
    query = """
    SELECT v.*, u.username as uploader 
    FROM videos v 
    JOIN users u ON v.uploaded_by = u.id
    """
    params = []
    
    if user_id:
        query += " WHERE v.uploaded_by = ?"
        params.append(user_id)
    
    query += " ORDER BY v.upload_time DESC"
    
    cur.execute(query, params)
    videos = [dict(row) for row in cur.fetchall()]
    conn.close()
    
    return videos

def save_trajectory(video_id, track_id, frame, bbox, confidence, timestamp=None):
    """保存轨迹数据"""
    conn = get_db_connection()
    cur = conn.cursor()
    
    x, y, width, height = bbox
    
    try:
        cur.execute(
            """
            INSERT INTO trajectories (video_id, track_id, frame, x, y, width, height, confidence, timestamp)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (video_id, track_id, frame, x, y, width, height, confidence, timestamp)
        )
        conn.commit()
        success = True
    except:
        success = False
    finally:
        conn.close()
    
    return success

def save_trajectories_batch(trajectories):
    """批量保存轨迹数据"""
    if not trajectories:
        return True
    
    conn = get_db_connection()
    cur = conn.cursor()
    
    try:
        cur.executemany(
            """
            INSERT INTO trajectories (video_id, track_id, frame, x, y, width, height, confidence, timestamp)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            trajectories
        )
        conn.commit()
        success = True
    except Exception as e:
        print(f"Error saving trajectories: {e}")
        success = False
    finally:
        conn.close()
    
    return success

def get_trajectories(video_id):
    """获取特定视频的轨迹数据"""
    conn = get_db_connection()
    cur = conn.cursor()
    
    cur.execute(
        """
        SELECT * FROM trajectories
        WHERE video_id = ?
        ORDER BY track_id, frame
        """,
        (video_id,)
    )
    
    trajectories = [dict(row) for row in cur.fetchall()]
    conn.close()
    
    return trajectories

def authenticate_user(username, password):
    """验证用户身份"""
    conn = get_db_connection()
    cur = conn.cursor()
    
    password_hash = hashlib.sha256(password.encode()).hexdigest()
    
    cur.execute(
        "SELECT id, username, is_admin FROM users WHERE username = ? AND password_hash = ?",
        (username, password_hash)
    )
    
    user = cur.fetchone()
    conn.close()
    
    if user:
        return dict(user)
    else:
        return None

# 若作为主程序运行，则初始化数据库
if __name__ == "__main__":
    init_db()
    print("数据库初始化完成。") 